/**
 * Copyright (c) 2022 Presto Labs Pte. Ltd.
 * Author: Ju
 */

#pragma once

#include <algorithm>
#include <atomic>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <memory>

#include "coin2/base/cpp_feature.h"

namespace coin2::base::memory {

//                   ----------- <--------- object
//                   | ref_cnt |
//                   |---------|
//                   | capacity|
//             |     |---------|<------|---- Buffer
//             |     |         |       |
//             |---->|---------|       |
//                   |         |       |
//                   |         |       |
//                   |         |       |
//                   |---------|<------|
//                   |         |
//                   -----------

class MemBuf : private base::NonCopyOrMovable {
 public:
  enum {
    kDefaultPacketSize = 2000,
  };

  // The following functions use malloc and free as the default allocator
  static MemBuf *Create(uint32_t size = kDefaultPacketSize);
  static MemBuf *Create(const void *buf, uint32_t size = kDefaultPacketSize);
  static MemBuf *Create(const void *buf, uint32_t size, uint32_t capacity);

  void AddRef() const;
  void Release() const;
  int32_t RefCount() const;
  const void *Begin() const;
  const void *End() const;

  void *Begin();
  void *End();

  MemBuf *Clone() const;
  uint32_t Capacity() const;

 private:
  explicit MemBuf(uint32_t size);
  MemBuf(const void *buf, uint32_t size);
  MemBuf(const void *buf, uint32_t size, uint32_t capacity);

  ~MemBuf() {}

  mutable std::atomic<int32_t> ref_cnt_;
  uint32_t capacity_;
};

struct MemBufDeleter {
  void operator()(MemBuf *p) const { p->Release(); }
};

typedef std::unique_ptr<MemBuf, MemBufDeleter> MemBufPtr;

inline int32_t MemBuf::RefCount() const {
  return ref_cnt_.load(std::memory_order_acquire);
}

inline void MemBuf::AddRef() const {
  ref_cnt_.fetch_add(1, std::memory_order_acq_rel);
}

inline void MemBuf::Release() const {
  if (ref_cnt_.fetch_sub(1, std::memory_order_acq_rel) <= 1) {
    free(const_cast<char *>(reinterpret_cast<const char *>(this)));
  }
}

inline const void *MemBuf::Begin() const {
  return reinterpret_cast<const void *>(this + 1);
}

inline const void *MemBuf::End() const {
  return reinterpret_cast<const char *>(this + 1) + capacity_;
}

inline void *MemBuf::Begin() { return reinterpret_cast<void *>(this + 1); }

inline void *MemBuf::End() {
  return reinterpret_cast<char *>(this + 1) + capacity_;
}

inline uint32_t MemBuf::Capacity() const { return capacity_; }

inline MemBuf *MemBuf::Create(uint32_t capacity) {
  auto p = malloc(sizeof(MemBuf) + capacity);
  return new (p) MemBuf(capacity);
}

inline MemBuf *MemBuf::Create(const void *buf, uint32_t capacity) {
  auto p = malloc(sizeof(MemBuf) + capacity);
  return new (p) MemBuf(buf, capacity);
}

inline MemBuf *MemBuf::Create(const void *buf, uint32_t size,
                              uint32_t capacity) {
  assert(size <= capacity);

  auto p = malloc(sizeof(MemBuf) + capacity);
  return new (p) MemBuf(buf, size, capacity);
}

inline MemBuf *MemBuf::Clone() const {
  auto p = MemBuf::Create(capacity_);
  ::memcpy(p->Begin(), Begin(), capacity_);
  return p;
}

inline MemBuf::MemBuf(const void *buf, uint32_t capacity) {
  ref_cnt_ = 0;
  capacity_ = capacity;
  ::memcpy(Begin(), buf, capacity);
}

inline MemBuf::MemBuf(const void *buf, uint32_t size, uint32_t capacity) {
  ref_cnt_ = 0;
  capacity_ = capacity;
  ::memcpy(Begin(), buf, (std::min)(size, capacity));
}

inline MemBuf::MemBuf(uint32_t capacity) {
  ref_cnt_ = 0;
  capacity_ = capacity;
}

}  // namespace coin2::base::memory